#include "vm.h"
#include "printf.h"
#include "console.h"
#include "string.h"

// NOTE: 本文件假设单核且恒等映射（PA==VA）。
// 如果将来采用高半区映射，只需修改 PA2VA/VA2PA 宏。

pagetable_t kernel_pagetable = 0;

extern char etext[]; // 链接脚本提供：内核 text 段结束

// ---------- create_pagetable ----------
pagetable_t create_pagetable(void) {
    pagetable_t pt = (pagetable_t)kalloc();
    if (pt == 0) {
        cprintf_color(31, "VM: failed to allocate root page table page\n");
        return 0;
    }
    // kalloc 返回页首地址（在恒等映射下同时为物理和虚拟地址）
    memset(pt, 0, PGSIZE);
    cprintf("VM: created new page table at %p\n", (void*)pt);
    return pt;
}

// ---------- walk_create ----------
// 遍历并在需要时创建中间页表页。返回指向 level0 PTE 的指针。
// 使用 PA2VA(PTE_PA(...)) 把 PTE 中的物理地址转成内核可解引用的指针。
pte_t* walk_create(pagetable_t pt, uint64_t va) {
    if (va >= MAXVA) {
        cprintf_color(31, "VM: virtual address 0x%lx exceeds MAXVA\n", va);
        return 0;
    }
    pagetable_t current_pt = pt;
    pte_t *pte;

    for (int level = 2; level > 0; level--) {
        int vpn_index = (int)VPN_MASK(va, level);
        pte = &current_pt[vpn_index];

        if (*pte & PTE_V) {
            // 有效项。若它是叶（有 R/W/X），直接返回该 PTE（superpage）
            if (*pte & (PTE_R | PTE_W | PTE_X)) {
                return pte;
            }
            // 否则取下一级页表指针
            uint64_t child_pa = PTE_PA(*pte);
            current_pt = (pagetable_t)PA2VA(child_pa);
        } else {
            // 分配新的页表页（作为下一级）
            pagetable_t new_pt = (pagetable_t)kalloc();
            if (new_pt == 0) {
                cprintf_color(31, "VM: failed to allocate level-%d page table\n", level);
                return 0;
            }
            memset(new_pt, 0, PGSIZE);
            // 把 PTE 指向新页表（设置 V 位，不设置 R/W/X）
            *pte = PA2PTE(new_pt) | PTE_V;
            current_pt = new_pt;
        }
    }

    // 返回最后一级 PTE 指针（level 0）
    int vpn_index = (int)VPN_MASK(va, 0);
    return &current_pt[vpn_index];
}

// ---------- walk_lookup ----------
pte_t* walk_lookup(pagetable_t pt, uint64_t va) {
    if (va >= MAXVA) return 0;
    pagetable_t current_pt = pt;
    pte_t *pte;

    for (int level = 2; level >= 0; level--) {
        int vpn_index = (int)VPN_MASK(va, level);
        pte = &current_pt[vpn_index];

        if (!(*pte & PTE_V)) {
            return 0; // 中间有无效项 -> 映射不存在
        }

        if (level > 0) {
            // 若是 leaf（R/W/X 任一置位），直接返回它（superpage）
            if (*pte & (PTE_R | PTE_W | PTE_X)) {
                return pte;
            }
            // 否则进入下一级
            uint64_t child_pa = PTE_PA(*pte);
            current_pt = (pagetable_t)PA2VA(child_pa);
        }
    }
    return pte;
}

// ---------- map_page ----------
int map_page(pagetable_t pt, uint64_t va, uint64_t pa, int perm) {
    if ((va & PGMASK) != 0 || (pa & PGMASK) != 0) {
        cprintf_color(31, "VM: unaligned address va=0x%lx pa=0x%lx\n", va, pa);
        return -1;
    }

    pte_t *pte = walk_create(pt, va);
    if (pte == 0) {
        cprintf_color(31, "VM: failed to walk page table for va=0x%lx\n", va);
        return -1;
    }

    if (*pte & PTE_V) {
        cprintf_color(33, "VM: remapping existing PTE for va=0x%lx\n", va);
        // 严格模式可直接返回错误
        // return -1;
    }

    *pte = PA2PTE(pa) | (perm & (PTE_R|PTE_W|PTE_X|PTE_U|PTE_G|PTE_A|PTE_D)) | PTE_V;
    // 在单核、恒等映射下不需要立即 sfence，但在真实硬件或多核应调用 sfence.vma
    return 0;
}

// ---------- map_region (按页逐页映射) ----------
int map_region(pagetable_t pt, uint64_t va, uint64_t pa, uint64_t size, int perm) {
    if (size == 0) return 0;
    uint64_t a = PGROUNDDOWN(va);
    uint64_t last = PGROUNDDOWN(va + size - 1);
    for (; ; a += PGSIZE, pa += PGSIZE) {
        if (map_page(pt, a, pa, perm) != 0) {
            cprintf_color(31, "VM: map_region failed at va=0x%lx pa=0x%lx\n", a, pa);
            return -1;
        }
        if (a == last) break;
    }
    return 0;
}

// ---------- destroy_pagetable_recursive ----------
static void destroy_pagetable_recursive(pagetable_t pt, int level) {
    if (pt == 0) return;
    if (level > 0) {
        for (int i = 0; i < 512; i++) {
            if (pt[i] & PTE_V) {
                // 若是 leaf（有 R/W/X），表示映射物理页，不应在此释放映射的物理页
                if (pt[i] & (PTE_R | PTE_W | PTE_X)) {
                    // leaf: 映射的物理页由调用者或其它逻辑释放（不在此处释放）
                } else {
                    uint64_t child_pa = PTE_PA(pt[i]);
                    pagetable_t child_pt = (pagetable_t)PA2VA(child_pa);
                    destroy_pagetable_recursive(child_pt, level - 1);
                }
            }
        }
    }
    // 释放当前页表页本身
    kfree((void*)pt);
}

// ---------- destroy_pagetable ----------
void destroy_pagetable(pagetable_t pt) {
    if (pt == 0) return;
    cprintf("VM: destroying page tree rooted at %p\n", (void*)pt);
    destroy_pagetable_recursive(pt, 2);
}

// ---------- dump_pagetable_recursive ----------
static void dump_pagetable_recursive(pagetable_t pt, int level, uint64_t base_va) {
    if (pt == 0) return;
    for (int i = 0; i < 512; i++) {
        if (pt[i] & PTE_V) {
            uint64_t current_va = base_va | ((uint64_t)i << VPN_SHIFT(level));
            // indent
            for (int j = 0; j < (2 - level) * 4; j++) consputc(' ');
            if (level == 0) {
                uint64_t pa = PTE_PA(pt[i]);
                // build perm string
                char perm[8] = {0};
                int p = 0;
                if (pt[i] & PTE_R) perm[p++] = 'R';
                if (pt[i] & PTE_W) perm[p++] = 'W';
                if (pt[i] & PTE_X) perm[p++] = 'X';
                if (pt[i] & PTE_U) perm[p++] = 'U';
                cprintf("VA:0x%lx -> PA:0x%lx [%s]\n", current_va, pa, perm);
            } else {
                uint64_t child_pa = PTE_PA(pt[i]);
                cprintf("L%d[%d] -> PT:%p\n", level, i, (void*)PA2VA(child_pa));
                dump_pagetable_recursive((pagetable_t)PA2VA(child_pa), level - 1, current_va);
            }
        }
    }
}

void dump_pagetable(pagetable_t pt, int level) {
    if (pt == 0) {
        cprintf("VM: cannot dump null page table\n");
        return;
    }
    cprintf("=== Page Table Dump (Root: %p) ===\n", (void*)pt);
    dump_pagetable_recursive(pt, level, 0);
    cprintf("=== End of Dump ===\n");
}

// ---------- SATP / TLB helpers ----------
static inline void w_satp(uint64_t x) {
    asm volatile("csrw satp, %0" : : "r"(x));
}
static inline void sfence_vma(void) {
    asm volatile("sfence.vma zero, zero");
}

// ---------- kvminit / kvminithart ----------
void kvminit(void) {
    kernel_pagetable = create_pagetable();
    if (!kernel_pagetable) {
        cprintf("KVMininit ERROR: cannot create kernel pagetable\n");
        while (1) { } // 停在这里，安全但不会崩溃
    }

    // Map kernel text (read + exec)
    uint64_t text_start = KERNBASE;
    uint64_t text_end = (uint64_t)etext;
    if (text_end < text_start) text_end = text_start; // safety
    map_region(kernel_pagetable, text_start, text_start, text_end - text_start, PTE_R | PTE_X);

    // Map kernel data + physical memory for allocation (read + write)
    uint64_t data_start = PGROUNDDOWN((uint64_t)etext);
    map_region(kernel_pagetable, data_start, data_start, PHYSTOP - data_start, PTE_R | PTE_W);

    // Map devices (UART0 as example)
    map_region(kernel_pagetable, UART0, UART0, PGSIZE, PTE_R | PTE_W);
    cprintf("KVMinitialized kernel pagetable root at %p\n", (void*)kernel_pagetable);
}

void kvminithart(void) {
    // kernel_pagetable currently is virtual pointer (identity). MAKE_SATP needs physical base.
    uint64_t satp_val = MAKE_SATP((uint64_t)kernel_pagetable);
    w_satp(satp_val);
    sfence_vma();
    cprintf("KVMinithart: satp set, paging enabled\n");
}
